#!/usr/bin/env python3
# Python implementation for IWR6843
# Data structure and reading of data loosely based on the matlab code from IWR6843 overhead people counting lab from Texas-Instruments
import rospy
from std_msgs.msg import Header
from sensor_msgs.msg import PointCloud2, PointField

from typing import Tuple, TextIO, BinaryIO
import platform
import serial
import struct
import binascii
import math as m
import argparse
import time
import sys
import json
import curses
import math
from copy import deepcopy
import numpy as np
import _thread

# Serial ports, can be overriden by specifying --dataport PORT and --controlport PORT
dataPort = '/dev/ttyUSB0'

# region: classess
class ULong: #uint64
    length: int = 8
    def __init__(self):
        pass
    
    def fromBytes(self, data: bytes, offset: int = 0):
        return struct.unpack_from('<Q', data, offset)[0]

class UInt: #uint32
    length: int = 4
    def __init__(self):
        pass
    
    def fromBytes(self, data: bytes, offset: int = 0):
        return struct.unpack_from('<I', data, offset)[0]

class Short: #int16
    length: int = 2
    def __init__(self):
        pass
    
    def fromBytes(self, data: bytes, offset: int = 0):
        return struct.unpack_from('<h', data, offset)[0]

class UShort: #uint16
    length: int = 2
    def __init__(self):
        pass
    
    def fromBytes(self, data: bytes, offset: int = 0):
        return struct.unpack_from('<H', data, offset)[0]

class UByte: #uint8
    length: int = 1
    def __init__(self):
        pass
    
    def fromBytes(self, data: bytes, offset: int = 0):
        return struct.unpack_from('<B', data, offset)[0]
        
class Byte: #int8
    length: int = 1
    def __init__(self):
        pass
    
    def fromBytes(self, data: bytes, offset: int = 0):
        return struct.unpack_from('<b', data, offset)[0]

class Float:
    length: int = 4
    def __init__(self):
        pass
    
    def fromBytes(self, data: bytes, offset: int = 0):
        return struct.unpack_from('<f', data, offset)[0]

#endregion

#region dataTypes

# Initiate the data types, used to convert from CTypes to Python Types
# To convert simply iterate over the values and increase the offset by the length property
frameHeaderStructType: dict = { #48 bytes long
    # 'sync':                         ULong(), # See syncPatternUINT64 below
    'version':                      UInt(),
    'totalPacketLen':               UInt(),
    'platform':                     UInt(),
    'frameNumber':                  UInt(), # Starting from 1
    'subframeNumber':               UInt(),
    'chirpProcessingMargin':        UInt(), # Chirp Processing margin, in ms
    'frameProcessingMargin':        UInt(), # Frame Processing margin, in ms
    'trackProcessTime':             UInt(), # Tracking Processing time, in ms
    'uartSentTime':                 UInt(), # Time spent to send data, in ms
    'numTLVs':                      UShort(), # Number of TLVs in thins frame
    'checksum':                     UShort() # Header checksum
}

frameHeaderLengthInBytes: int = 48
tlvHeaderLengthInBytes: int = 8
pointLengthInBytes: int = 8
pointUnitInBytes: int = 20
targetLengthInBytes: int = 40

tlvHeaderStruct: dict = { # 8 bytes
    'type':             UInt(), # TLV object Type
    'length':           UInt() # TLV object Length, in bytes, including TLV header
}


# Point Cloud TLV object consists of an array of points. 
# Each point has a structure defined below
# Type 6
pointStruct: dict = { # 20 bytes
    'elevation':             Byte(), 
    'azimuth':            Byte(), # Angel, in rad
    'doppler':          Short(), # Doplper, in m/s
    'range':            UShort(), # Range, in m
    'snr':              UShort()    # SNR, ratio
}

pointUnit: dict = {
    'elevationUnit':    Float(),
    'azimuthUnit':      Float(),
    'dopplerUnit':      Float(),
    'rangeUnit':        Float(),
    'snrUnit':          Float(),
}
# Target List TLV object consists of an array of targets. 
# Each target has a structure define below
# Type 7
targetStruct: dict = { #40 bytes
    'tid':              UInt(),# Track ID
    'posX':             Float(), # Target position in X dimension, m
    'posY':             Float(), # Target position in Y dimension, m
    'posZ':             Float(),
    'velX':             Float(), # Target velocity in X dimension, m/s
    'velY':             Float(), # Target velocity in Y dimension, m/s
    'velZ':             Float(),
    'accX':             Float(), # Target acceleration in X dimension, m/s2
    'accY':             Float(), # Target acceleration in Y dimension, m/s
    'accZ':             Float()
}

cTypesInfo: dict = {
    'uint64': "<Q",
    'uint32': "<I",
    'uint16': "<H",
    'uint8': "<B"
}

#endregion


serialData = serial.Serial()
serialData.baudrate = 921600
serialData.port = dataPort
serialData.open()

syncPattern = b'\x02\x01\x04\x03\x06\x05\x08\x07'

# Simply counter to keep track of number of frames captured
dataCount = 0
targetData: dict = dict()

def polar_to_carte(ele, azi, radius):
    x = np.cos(ele) * np.sin(azi) * radius
    y = np.cos(ele) * np.cos(azi) * radius
    z = np.sin(ele) * radius
    return x, y, z

def transform_coord(reflection):
    reflection = deepcopy(reflection)
    reflection[:, 0], reflection[:, 1], reflection[:, 3] = \
        polar_to_carte(reflection[:, 0], reflection[:, 1], reflection[:, 3])
    
    return reflection

# This function keeps a buffer of 8 bytes and compares it with the sync pattern
# When sync pattern is found the function returns to the caller's context
def WaitForSyncPattern():
    # Initialize with the syncPattern
    # since data gets popped before the first compare this never causes it to be true but does ensure the length is the same
    dataBuffer = bytearray(syncPattern)

    # Wait for the serial port to receive the sync pattern
    while True:
        # Listen for ctrl+c in case it gets stuck
        # Read 1 byte of data
        currentData = serialData.read(1)
        # Remove the oldest byte received from the serial port
        dataBuffer.pop(0)
        dataBuffer += currentData
        # If binfile is set (dump file parameter on command line) write raw data to file
        
        # Break the loop if the data is the same as the syncPattern
        if bytes(dataBuffer) == syncPattern:
            return

rospy.init_node('iwr6843_driver', anonymous=True)
rate = rospy.Rate(10) # 10hz
pub = rospy.Publisher('/iwr6843_points', PointCloud2, queue_size=10)

while not rospy.is_shutdown():
    numTargets = 0
    numPoints = 0
    seq_cnt = 0
    WaitForSyncPattern()

    # Initialize the data dictionary
    dataDict = {
        'tlvHeaders': list()
    }
    # offset in bytes in the current packet, to ensure the correct bytes are read
    offset: int = 0

    # Decode the header, iterate over every type as defined in frameHeaderStructType and decode it to the corresponding data type
    # Same process is repeated for every other type of data (tlvHeader, pointCloud and target)
    for dataPoint in frameHeaderStructType.items():
        currentData = serialData.read(dataPoint[1].length)
        dataDict[dataPoint[0]] = dataPoint[1].fromBytes(currentData, 0)

    # length of only the data in bytes, so just the packet without the header, used for reading the whole packet in one go
    dataLength = dataDict['totalPacketLen'] - frameHeaderLengthInBytes
    # Read the remaining data
    data = serialData.read(dataLength)
    # print("needs to Read ", dataLength, " Bytes from serial", ", actual data size ", sys.getsizeof(data), " TLV numbers:", dataDict['numTLVs'])

    for nTlv in range(0, dataDict['numTLVs']):
        tlvData = {}
        # Decode TLV Header
        for dataPoint in tlvHeaderStruct.items():
            tlvData[dataPoint[0]] = dataPoint[1].fromBytes(data, offset)
            offset += dataPoint[1].length

        # Length of the TLV Data in bytes, this is the total length as specified in the tlvHeader minus the tlvHeader
        valueLength = tlvData['length'] - tlvHeaderLengthInBytes
        # print("TLV ", nTlv, "carried length value : ", tlvData['length'])
        # Decode TLV Data
        if tlvData['type'] == 6: # type 6: Point Cloud Data
            # Number of points, so the data length divided by the length in bytes of one point cloud data point
            numPoints = m.floor((valueLength - pointUnitInBytes)/pointLengthInBytes)
            point_unitvalue = []
            for i, pu in enumerate(pointUnit.items()):
                point_unitvalue.append(pu[1].fromBytes(data, offset))
                offset += pu[1].length

            radar_pcl = PointCloud2()
            seq_cnt +=1
            radar_pcl.header.seq = seq_cnt
            radar_pcl.header.stamp = rospy.Time.now()
            radar_pcl.header.frame_id = "map"
            radar_pcl.point_step = 20 # 4 bytes for each value, 5 values in total ()
            radar_pcl.row_step = numPoints * 20 

            radar_pcl.height = 1
            
            
            radar_pcl_field1 = PointField()
            radar_pcl_field1.name = "x"
            radar_pcl_field1.datatype = 7
            radar_pcl_field1.offset = 0
            radar_pcl_field1.count = 1
            radar_pcl.fields.append(radar_pcl_field1)

            radar_pcl_field2 = PointField()
            radar_pcl_field2.name = "y"
            radar_pcl_field2.datatype = 7
            radar_pcl_field2.offset = 4
            radar_pcl_field2.count = 1
            radar_pcl.fields.append(radar_pcl_field2)

            radar_pcl_field3 = PointField()
            radar_pcl_field3.name = "v"
            radar_pcl_field3.datatype = 7
            radar_pcl_field3.offset = 8
            radar_pcl_field3.count = 1
            radar_pcl.fields.append(radar_pcl_field3)

            radar_pcl_field4 = PointField()
            radar_pcl_field4.name = "z"
            radar_pcl_field4.datatype = 7
            radar_pcl_field4.offset = 12
            radar_pcl_field4.count = 1
            radar_pcl.fields.append(radar_pcl_field4)

            radar_pcl_field5 = PointField()
            radar_pcl_field5.name = "snr"
            radar_pcl_field5.datatype = 7
            radar_pcl_field5.offset = 16
            radar_pcl_field5.count = 1
            radar_pcl.fields.append(radar_pcl_field5)

            points = []
            for i in range(0, numPoints):
                point = []
                # Could add data parsing here, but not necessary since we don't need the pointcloud
                for i, p in enumerate(pointStruct.items()):
                    point.append(p[1].fromBytes(data, offset) * point_unitvalue[i])
                    offset += p[1].length
                points.append(point)
            points = np.array(points)
            points = transform_coord(points)
            saved = False
            if not saved:
                np.save("pt.npy", points)
                saved = True
            points = points[points[:,2]!=0]
            radar_pcl.data = points.flatten().astype(np.float32).tobytes()
            # radar_pcl.data = points
            numPoints = points.shape[0]
            radar_pcl.width = numPoints
            radar_pcl.row_step = numPoints * 20 
            pub.publish(radar_pcl)
                # print(point)
    dataCount += 1


# Close the serial connections
serialData.close()
